//  原版权信息 LAC v3

//  修改者：
// Copyright (c) 2025 刘颖，张宇
//  所属单位：中国科学院信息工程研究所网络空间安全防御重点实验室
//             Key Laboratory of Cyberspace Security Defense,Institute of Information Engineering, CAS
//             中国科学院大学网络空间安全学院
//             School of Cyber Security, University of Chinese Academy of Sciences

// Performance and Optimization Features:
// - Compact lifted modulus NTT-based multiplication, implemented in poly_mul and poly_aff functions [Ying Liu]
// - Sparse ternary distribution sampling, implemented in gen_e function [Yu Zhang]


#include <immintrin.h>
#include <string.h>
#include "bin-lwe.h"
#include "rand.h"
#include "polarlac_param.h"
#include "ntt.h"
#include "ntt-1024.h"
#include "fips202.h"

#include <stdio.h>

#ifdef LINUX
#define ALIGNED(x) __attribute__((aligned(x)))
#endif

#ifdef WIN
#define ALIGNED(x) __declspec(align(x))
#endif

static unsigned int rej_uniform(uint8_t *r, unsigned int len, const uint8_t *buf, unsigned int buflen)
{
  unsigned int ctr, pos;

  ctr = pos = 0;
  while(ctr < len && pos<= buflen) 
  {
    if(buf[pos] < Q)
      r[ctr++] = buf[pos];
	pos++;
  }

  return ctr;
}

// generate the public parameter a from seed
int gen_a(unsigned char *a, const unsigned char *seed)
{
	unsigned int ctr, buflen;
	keccak_state state;
	shake256_absorb_once(&state, seed, SEED_LEN);

#if defined LAC256
	uint8_t buf[1088];
	shake256_squeezeblocks(buf, 8, &state);
	shake256_squeeze(buf + 408, 44, &state);
	buflen = 1088;
#else
	uint8_t buf[544];
	shake256_squeezeblocks(buf, 4, &state);
	buflen = 544;
#endif

	ctr = rej_uniform(a, DIM_N, buf, buflen);
	while(ctr < DIM_N)
	{
		shake256_squeezeblocks(buf, 1, &state);
		buflen = SHAKE256_RATE;
		ctr += rej_uniform(a + ctr, DIM_N - ctr, buf, buflen);
	}

	return 0;
}

// The generation of errors with constant time.
// -1 will be stord as 255 with the unsigned char form
int gen_e(unsigned char *e, keccak_state *state)
{
	uint8_t flag;
	uint8_t r[SHAKE256_RATE * 2], tmp[DIM_N * 2];
	uint16_t i, j, t;
	uint16_t mask, norm;
	uint16_t e_1,e_2;
	memset(e, 0, DIM_N);
#if defined LAC_LIGHT
	shake256_squeezeblocks(r, 1, state);
	t = 0;
	for (i = 0; i < 8; i++)
	{
		for (j = 0; j < DIM_N_8 * 2; j++)
		{
			tmp[i * (DIM_N_8 * 2) + j] = (r[t + j] & 1);
			r[t + j] = (r[t + j] >> 1);
		}
	}
	for (i = 0; i < DIM_N; i++)
	{
		tmp[i] = tmp[i] - tmp[i + DIM_N];
	}

	flag = 1;
	while(flag)
	{
		shake256_squeeze(r, DIM_N_8, state);
		t = 0;
		for (i = 0; i < 8; i++)
		{
			for (j = 0; j < DIM_N_8; j++)
			{
				tmp[DIM_N + i * DIM_N_8 + j] = (r[t + j] & 1);
				r[t + j] = (r[t + j] >> 1);
			}
		}

		flag = 0;
		norm = 0;
		for(i = 0; i < DIM_N; i++)
		{
			e[i] = tmp[i] * tmp[i + DIM_N];
			norm += (e[i] & e[i] & 1);
		}
		mask = (norm < 110);
		flag = (1 & (-mask)) | (flag & (~(-mask)));
		mask = (norm > 146);
		flag = (1 & (-mask)) | (flag & (~(-mask)));

	}
	
	for (i = 0; i < DIM_N; i++)
	{
		// e[i] = tmp[i] * tmp[i + DIM_N];
		mask = (e[i] > q_half);
		e_1 = neg_one;
		e_2 = e[i];
		e[i] = (e_1 & (-mask)) | (e_2 & (~(-mask)));
	}

#elif defined LAC128 // 1/2 Hamming Weight（b0-b1)
	shake256_squeeze(r, DIM_N_8, state);
	t = 0;
	for (i = 0; i < 8; i++)
	{
		for (j = 0; j < DIM_N_8; j++)
		{
			tmp[i * DIM_N_8 + j] = (r[t + j] & 1);
			r[t + j] = (r[t + j] >> 1);
		}
	}

	flag = 1;
	while(flag)
	{
		shake256_squeeze(r, DIM_N_8, state);
		t = 0;
		for (i = 0; i < 8; i++)
		{
			for (j = 0; j < DIM_N_8; j++)
			{
				tmp[i * DIM_N_8 + j + DIM_N] = (r[t + j] & 1);
				r[t + j] = (r[t + j] >> 1);
			}
		}
		flag = 0;
		norm = 0;
		for(i = 0; i < DIM_N; i++)
		{
			e[i] = tmp[i] - tmp[i + DIM_N];
			norm += (e[i] & e[i] & 1);
		}
		mask = (norm < 248);
		flag = (1 & (-mask)) | (flag & (~(-mask)));
		mask = (norm > 264);
		flag = (1 & (-mask)) | (flag & (~(-mask)));
	}

	for (i = 0; i < DIM_N; i++)
	{
		// e[i] = tmp[i] - tmp[i + DIM_N];
		//trans 255 to 250
		mask = (e[i] > q_half);
		e_1 = neg_one;
		e_2 = e[i];
		e[i] = (e_1 & (-mask)) | (e_2 & (~(-mask)));
	}

#else // (b0-b1)*b2
	shake256_squeezeblocks(r, 2, state);
	t = 0;
	for (i = 0; i < 8; i++)
	{
		for (j = 0; j < DIM_N_8 * 2; j++)
		{
			tmp[i * (DIM_N_8 * 2) + j] = (r[t + j] & 1);
			r[t + j] = (r[t + j] >> 1);
		}
	}
	for (i = 0; i < DIM_N; i++)
	{
		tmp[i] = tmp[i] - tmp[i + DIM_N];
	}
	
	flag = 1;
	while(flag)
	{
		shake256_squeezeblocks(r, 1, state);
		t = 0;
		for (i = 0; i < 8; i++)
		{
			for (j = 0; j < DIM_N_8; j++)
			{
				tmp[i * DIM_N_8 + j + DIM_N] = (r[t + j] & 1);
				r[t + j] = (r[t + j] >> 1);
			}
		}
		flag = 0;
		norm = 0;
		for(i = 0; i < DIM_N; i++)
		{
			e[i] = tmp[i] * tmp[i + DIM_N];
			norm += (e[i] & e[i] & 1);
		}
		mask = (norm < 242);
		flag = (1 & (-mask)) | (flag & (~(-mask)));
		mask = (norm > 270);
		flag = (1 & (-mask)) | (flag & (~(-mask)));
	}

	for (i = 0; i < DIM_N; i++)
	{
		// e[i] = tmp[i] * tmp[i + DIM_N];
		mask = (e[i] > q_half);
		e_1 = neg_one;
		e_2 = e[i];
		e[i] = (e_1 & (-mask)) | (e_2 & (~(-mask)));
	}
#endif

	return 0;
}

// b=as using compact lift multiplication with constant time.
int poly_mul(const unsigned char *a, const unsigned char *s, unsigned char *b, unsigned int vec_num)
{
	int i;

	uint16_t a2[DIM_N], s2[DIM_N], b2[DIM_N];
	uint16_t mask;
	uint16_t a_1, a_2, s_1, s_2, b_1, b_2;

	// step 1: map to the lifted ring with NTTQ as the modulus
	for (i = 0; i < DIM_N; i++)
	{
		mask = (a[i] >= q_half);
		a_1 = a[i] + Q_sub_q;
		a_2 = a[i] + NTTQ;
		a2[i] = (a_1 & (-mask)) | (a_2 & (~(-mask)));
	}

	for (i = 0; i < DIM_N; i++)
	{
		mask = (s[i] >= q_half);
		// s_1 = s[i] + Q_sub_uint8;
		s_1 = s[i] + Q_sub_q;
		s_2 = s[i] + NTTQ;
		s2[i] = (s_1 & (-mask)) | (s_2 & (~(-mask)));
	}

#if defined LAC_LIGHT || defined LAC128
	// step 2: perform ntt-based multiplciation with NTTQ as the modulus
	poly_mul_ntt_Lazy(a2, s2, b2);
#else 
	poly_mul_ntt_Lazy_1024(a2, s2, b2);
#endif
	// step 3: map back to the original ring with Q as the modulus
	for (i = 0; i < vec_num; i++)
	{
		mask = (b2[i] < Q_half);
		b_1 = b2[i] + NTTQ;
		b_2 = b2[i];
		b2[i] = (b_1 & (-mask)) | (b_2 & (~(-mask)));
		b[i] = (b2[i] + neg_Q_mod_q) % Q;
	}

	return 0;
}

// b=as+e using compact lift multiplication with constant time.
int poly_aff(const unsigned char *a, const unsigned char *s, unsigned char *e, unsigned char *b, unsigned int vec_num)
{
	int i;

	uint16_t a2[DIM_N], s2[DIM_N], b2[DIM_N];

	// test ntt
	uint16_t mask;
	uint16_t a_1, a_2, s_1, s_2, b_1, b_2;
	for (i = 0; i < DIM_N; i++)
	{
		mask = (a[i] >= q_half);
		a_1 = a[i] + Q_sub_q;
		a_2 = a[i] + NTTQ;
		a2[i] = (a_1 & (-mask)) | (a_2 & (~(-mask)));
	}

	for (i = 0; i < DIM_N; i++)
	{
		mask = (s[i] >= q_half);
		s_1 = s[i] + Q_sub_q;
		s_2 = s[i] + NTTQ;
		s2[i] = (s_1 & (-mask)) | (s_2 & (~(-mask)));
	}

#if defined LAC_LIGHT || defined LAC128
	poly_mul_ntt_Lazy(a2, s2, b2);
#else
	poly_mul_ntt_Lazy_1024(a2, s2, b2);
#endif

	for (i = 0; i < vec_num; i++)
	{
		mask = (b2[i] < Q_half);
		b_1 = b2[i] + NTTQ;
		b_2 = b2[i];
		b2[i] = (b_1 & (-mask)) | (b_2 & (~(-mask)));
		b2[i] = (b2[i] + neg_Q_mod_q) % Q;
		b[i] = (b2[i] + e[i] + Q) % Q;
	}

	return 0;
}

// Compression: c1 discards 1-bit
int poly_compress_c1_1bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,j,loop;
	loop = vec_num / 8;
    for (i = 0; i < loop; i++) {
        unsigned char buf[8];
        for (j = 0; j < 8; j++) {
            buf[j] = (in[i * 8 + j]) >> 1;
        }

        out[i * 7 + 0] = (buf[1] << 7) | buf[0];
        out[i * 7 + 1] = (buf[2] << 6) | (buf[1] >> 1);
        out[i * 7 + 2] = (buf[3] << 5) | (buf[2] >> 2);
        out[i * 7 + 3] = (buf[4] << 4) | (buf[3] >> 3);
        out[i * 7 + 4] = (buf[5] << 3) | (buf[4] >> 4);
        out[i * 7 + 5] = (buf[6] << 2) | (buf[5] >> 5);
        out[i * 7 + 6] = (buf[7] << 1) | (buf[6] >> 6);
    }
	
	return 0;
}
int poly_decompress_c1_1bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,loop;
	loop = vec_num / 8;
    for (i = 0; i < loop; i++) {
        out[i * 8 + 0] = (in[i * 7 + 0] << 1) + 0b00000001;
		out[i * 8 + 1] = ((in[i * 7 + 1] << 2) | ((in[i * 7 + 0] & 0b10000000) >> 6)) + 0b00000001;
		out[i * 8 + 2] = ((in[i * 7 + 2] << 3) | ((in[i * 7 + 1] & 0b11000000) >> 5)) + 0b00000001;
		out[i * 8 + 3] = ((in[i * 7 + 3] << 4) | ((in[i * 7 + 2] & 0b11100000) >> 4)) + 0b00000001;
		out[i * 8 + 4] = ((in[i * 7 + 4] << 5) | ((in[i * 7 + 3] & 0b11110000) >> 3)) + 0b00000001;
		out[i * 8 + 5] = ((in[i * 7 + 5] << 6) | ((in[i * 7 + 4] & 0b11111000) >> 2)) + 0b00000001;
		out[i * 8 + 6] = ((in[i * 7 + 6] << 7) | ((in[i * 7 + 5] & 0b11111100) >> 1)) + 0b00000001;
		out[i * 8 + 7] = (in[i * 7 + 6] & 0b11111110) + 0b00000001;
    }
	
	return 0;
}

// Compression: c2 discards 4-bit
int poly_compress_c2_4bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,loop;
	loop=vec_num/2;
	for(i=0;i<loop;i++)
	{
		out[i]=(in[i*2])>>4;
		out[i]=out[i]^(in[i*2+1]&0xf0);
	}
	
	return 0;
}
int poly_decompress_c2_4bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,loop;
	loop=vec_num/2;
	for(i=0;i<loop;i++)
	{
		out[i*2]=(in[i]<<4)^0x08;
		out[i*2+1]=(in[i]&0xf0)^0x08;
	}
	
	return 0;
}

// Compression: c2 discards 5-bit
int poly_compress_c2_5bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,j,loop;
	loop = vec_num / 8;
    for (i = 0; i < loop; i++) {
        unsigned char buf[8];
        for (j = 0; j < 8; j++) {
            buf[j] = in[i * 8 + j] >> 5;
        }

        out[i * 3 + 0] = (buf[2] << 6) | (buf[1] << 3) | buf[0];
        out[i * 3 + 1] = (buf[5] << 7) | (buf[4] << 4) | (buf[3] << 1) | (buf[2] >> 2);
        out[i * 3 + 2] = (buf[7] << 5) | (buf[6] << 2) | (buf[5] >> 1);
    }
	
	return 0;
}
int poly_decompress_c2_5bit(const unsigned char *in,  unsigned char *out, const unsigned int vec_num)
{
	int i,loop;
	loop = vec_num / 8;
    for (i = 0; i < loop; i++) {
        out[i * 8 + 0] = (in[i * 3 + 0] << 5) + 0b00010000;
		out[i * 8 + 1] = ((in[i * 3 + 0] & 0b00111000) << 2) + 0b00010000;
		out[i * 8 + 2] = ((in[i * 3 + 1] << 7) | (in[i * 3 + 0] & 0b11000000) >> 1) + 0b00010000;
		out[i * 8 + 3] = ((in[i * 3 + 1] & 0b00001110) << 4) + 0b00010000;
		out[i * 8 + 4] = ((in[i * 3 + 1] & 0b01110000) << 1) + 0b00010000;
		out[i * 8 + 5] = ((in[i * 3 + 2] << 6) | (in[i * 3 + 1] & 0b10000000) >> 2) + 0b00010000;
		out[i * 8 + 6] = ((in[i * 3 + 2] & 0b00011100) << 3) + 0b00010000;
		out[i * 8 + 7] = (in[i * 3 + 2] & 0b11100000) + 0b00010000;
    }
	
	return 0;
}
